

use std::thread::spawn;

use clap::Parser;
use serde::{Deserialize, Serialize};

use wtask_base::{
    error::{WError, WResult},
    export::{KillReceiver, TcpStream, log_debug, log_info, select},
    machine::{CMDArgsConfigMachine, WSpawn, runtime_init},
    tcp::{WTcpStream, stream_exchange_akey, tcp_listen}
};

use crate::{basic::Empty, func::cli_app_stream, utils::{AppTrait, StreamStatus}};




#[derive(Debug, Clone, Serialize, Deserialize, Parser)]
pub struct AppPort {
    /// 远程端口
    pub port_remote: u32,

    /// 本地端口，如果不提供则使用远程的端口
    pub port_local: Option<u32>,

    #[arg(long, default_value = "1024")]
    pub capacity: usize,
}

impl AppPort {
    async fn port_session(self, stream_session: WTcpStream, machine: CMDArgsConfigMachine) -> WResult<()> {
        let capacity = self.capacity;
        let (stream, key) = cli_app_stream(self, machine).await?;
        stream_exchange_akey(&stream, &stream_session, &key, capacity).await?;
        Ok(())
    }
    async fn server_conn_forwarding(&self,  stream_session: &WTcpStream, session_token: &str, num: u32) -> WResult<()> {
        log_debug!("App port get a session: {session_token} {num}, Connect forwarding port ...");
        let stream = TcpStream::connect(format!("127.0.0.1:{}", self.port_remote)).await?;
        log_info!("Machine Connect 127.0.0.1:{}, port-forward", self.port_remote);
        let stream_normal = stream.into_split();
        log_debug!("stream port exchange data: {session_token}");
        let _ = stream_exchange_akey(stream_session, &stream_normal, session_token, self.capacity).await;
        Ok(())
    }
}


impl AppTrait for AppPort {
    type Status = (KillReceiver<bool>, CMDArgsConfigMachine);
    type Res = Empty;

    fn client_handle_before<M: FnMut(bool, String) + Send + Sync>(self, _msg_func: M, status: Option<Self::Status>) -> WResult<(Option<Self>, Option<Self::Status>)> {
        let status = status.ok_or(WError::Empty)?;
        let mut kill_rx_parent = status.0;
        let port = if let Some(port) = self.port_local {
            port
        } else {
            self.port_remote
        };
        let machine = status.1;
        let handle: std::thread::JoinHandle<Result<(), WError>> = spawn(move || {
            let runtime  = runtime_init(8)?;
            let _t: Result<(), WError> = runtime.block_on(async move {
                let listener = tcp_listen(port, "PortForwarding")?;
                let mut wspawn = WSpawn::default();
                let (kill_tx, mut kill_rx) = wspawn.catch_pause_signal(async move {});
                let port_info = format!("{} {} => 0.0.0.0:{}", machine, self.port_remote, port);
                log_info!("PortForwarding: {}", port_info);
                loop {
                    let app = self.clone();
                    let machine = machine.clone();
                    select! {
                        Ok((stream, _)) = listener.accept() => {
                            // log_info!("Get a port-forwarding connection: {} ...", port_info);
                            wspawn.spawn(async move {
                                let stream_now = stream.into_split();
                                let v = app.port_session(stream_now, machine).await;
                                if let Err(e) = v {
                                    log_debug!("session port forwarding err stop: {e:?}");
                                }
                            });
                        },
                        _ = kill_rx.changed() => {
                            log_info!("Shutting PortForwarding: {} ...", port_info);
                            break;
                        },
                        _ = kill_rx_parent.changed() => {
                            log_debug!("Shutting PortForwarding {}, kill by parent ...", port_info);
                            let _ = kill_tx.send(true);
                            break;
                        },
                    }
                };
                Ok(())
            });
            runtime.shutdown_background();
            Ok(())
        });
        handle.join().unwrap()?;
        Ok((None, None))
    }

    async fn server_handle(self, _server_token: &str, session_token: &str, stream_session: &WTcpStream, stream_status: StreamStatus) -> WResult<()> {
        let is_old = stream_status.read().await.2.get(&self.port_remote).cloned();
        if let Some(n) = is_old {
            stream_status.write().await.2.insert(self.port_remote, n+1);
            self.server_conn_forwarding(stream_session, session_token, n+1).await?;
        } else {
            stream_status.write().await.2.insert(self.port_remote, 1);
            self.server_conn_forwarding(stream_session, session_token, 1).await?;
        };
        let n = stream_status.read().await.2.get(&self.port_remote).cloned().unwrap();
        if n == 1 {
            stream_status.write().await.2.remove(&self.port_remote);
        } else {
            stream_status.write().await.2.insert(self.port_remote, n-1);
        }
        Ok(())
    }
}

